module maxpool_2x2 #(
    parameter int unsigned P_CH  = 4,
    parameter int unsigned N_CH  = 16,
    parameter int unsigned N_IH  = 64,
    parameter int unsigned N_IW  = 64,
    parameter int unsigned A_BIT = 8
) (
    input  logic                  clk,
    input  logic                  rst_n,
    input  logic [P_CH*A_BIT-1:0] in_data,
    input  logic                  in_valid,
    output logic                  in_ready,
    output logic [P_CH*A_BIT-1:0] out_data,
    output logic                  out_valid,
    input  logic                  out_ready
);

    initial begin
        if (N_CH < P_CH) begin
            $fatal(1, "maxpool_2x2: N_CH must be >= P_CH");
        end
        if (N_CH % P_CH != 0) begin
            $fatal(1, "maxpool_2x2: N_CH must be divisible by P_CH");
        end
        if (N_IH % 2 != 0) begin
            $fatal(1, "maxpool_2x2: N_IH must be even");
        end
        if (N_IW % 2 != 0) begin
            $fatal(1, "maxpool_2x2: N_IW must be even");
        end
    end

    localparam int unsigned FOLD = N_CH / P_CH;
    localparam int unsigned N_OH = N_IH / 2;
    localparam int unsigned N_OW = N_IW / 2;

    logic [    P_CH*A_BIT-1:0] line_buf    [N_OW][FOLD];
    logic [$clog2(N_IH+1)-1:0] cntr_h;
    logic [$clog2(N_IW+1)-1:0] cntr_w;
    logic [$clog2(FOLD+1)-1:0] cntr_f;
    logic                      pipe_en;
    logic                      is_row_even;
    logic                      is_col_even;
    logic                      is_lst_fold;
    logic [$clog2(N_OW+1)-1:0] idx;

    assign pipe_en     = in_valid && out_ready;
    assign in_ready    = out_ready;
    assign is_row_even = (cntr_h[0] == 1'b0);
    assign is_col_even = (cntr_w[0] == 1'b0);
    assign is_lst_fold = (cntr_f == FOLD - 1);
    assign idx         = cntr_w >> 1;

    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            cntr_h <= 0;
            cntr_w <= 0;
            cntr_f <= 0;
        end else if (pipe_en) begin
            if (cntr_f == FOLD - 1) begin
                cntr_f <= 0;
                if (cntr_w == N_IW - 1) begin
                    cntr_w <= 0;
                    if (cntr_h == N_IH - 1) begin
                        cntr_h <= 0;
                    end else begin
                        cntr_h <= cntr_h + 1;
                    end
                end else begin
                    cntr_w <= cntr_w + 1;
                end
            end else begin
                cntr_f <= cntr_f + 1;
            end
        end
    end

    function automatic logic [P_CH*A_BIT-1:0] max_vec(input logic [P_CH*A_BIT-1:0] a, input logic [P_CH*A_BIT-1:0] b);
        logic [A_BIT-1:0] a_ch, b_ch;
        for (int i = 0; i < P_CH; i++) begin
            a_ch                    = a[i*A_BIT+:A_BIT];
            b_ch                    = b[i*A_BIT+:A_BIT];
            max_vec[i*A_BIT+:A_BIT] = (a_ch > b_ch) ? a_ch : b_ch;
        end
    endfunction

    logic [P_CH*A_BIT-1:0] out_buf_d1;
    logic                  out_valid_d1;
    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            for (int i = 0; i < N_OW; i++) begin
                for (int f = 0; f < FOLD; f++) begin
                    line_buf[i][f] <= '0;
                end
            end
            out_buf_d1   <= '0;
            out_valid_d1 <= 1'b0;
        end else if (pipe_en) begin
            out_valid_d1 <= 1'b0;
            if (is_row_even) begin
                if (is_col_even) begin
                    line_buf[idx][cntr_f] <= in_data;
                end else begin
                    line_buf[idx][cntr_f] <= max_vec(in_data, line_buf[idx][cntr_f]);
                end
            end else begin
                if (is_col_even) begin
                    line_buf[idx][cntr_f] <= max_vec(in_data, line_buf[idx][cntr_f]);
                end else begin
                    out_buf_d1   <= max_vec(in_data, line_buf[idx][cntr_f]);
                    out_valid_d1 <= 1'b1;
                end
            end
        end else if (out_ready) begin
            out_valid_d1 <= 1'b0;
        end
    end

    assign out_data  = out_buf_d1;
    assign out_valid = out_valid_d1;

endmodule
